package cn.itcast.hello;

import org.apache.spark.ml.evaluation.RegressionEvaluator;
import org.apache.spark.ml.feature.VectorAssembler;
import org.apache.spark.ml.regression.LinearRegression;
import org.apache.spark.ml.regression.LinearRegressionModel;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.streaming.StreamingQueryException;

import java.io.IOException;
import java.util.concurrent.TimeoutException;

/**
 * Author itcast
 * Desc 演示使用Java语言开发SparkMlLib-线性回归算法-房价预测案例
 */
public class JavaSparkDemo05 {
    public static void main(String[] args) throws TimeoutException, StreamingQueryException, IOException {
        //0.TODO 准备环境
        SparkSession spark = SparkSession.builder().appName("JavaSparkDemo").master("local[*]")
                .config("spark.sql.shuffle.partitions", "4")
                .getOrCreate();
        spark.sparkContext().setLogLevel("WARN");

        //TODO 1.加载数据
        Dataset<Row> homeDataDF = spark.read()
                .format("csv")
                .option("sep", "|")//指定分隔符
                .option("header", "true")//是否有表头
                .option("inferSchema", "true")//是否自动推断约束
                .load("data/input/homeprice.data");
        homeDataDF.printSchema();
        homeDataDF.show();
        /*
        root
 |-- mlsNum: integer (nullable = true)
 |-- city: string (nullable = true)
 |-- sqFt: double (nullable = true)
 |-- bedrooms: integer (nullable = true)
 |-- bathrooms: integer (nullable = true)
 |-- garage: integer (nullable = true)
 |-- age: integer (nullable = true)
 |-- acres: double (nullable = true)
 |-- price: double (nullable = true)
//|房屋编号|城市|平方英尺|卧室数|卫生间数|车库|年龄|房屋占地面积|房屋价格
+-------+------------+-------+--------+---------+------+---+-----+---------+
| mlsNum|        city|   sqFt|bedrooms|bathrooms|garage|age|acres|    price|
+-------+------------+-------+--------+---------+------+---+-----+---------+
|4424109|Apple Valley| 1634.0|       2|        2|     2| 33| 0.04| 119900.0|
|4404211|   Rosemount|13837.0|       4|        6|     4| 17|14.46|3500000.0|
|4339082|  Burnsville| 9040.0|       4|        6|     8| 12| 0.74|2690000.0|
         */

        //TODO 2.特征处理
        //特征选择
        Dataset<Row> featuredDF = homeDataDF.select("sqFt", "age", "acres", "price");
        //特征向量化
        Dataset<Row> vectorDF = new VectorAssembler()
                .setInputCols(new String[]{"sqFt", "age", "acres"})//指定要对哪些特征做向量化
                .setOutputCol("features")//向量化之后的特征列列名
                .transform(featuredDF);
        vectorDF.printSchema();
        vectorDF.show();
        /*
        root
         |-- sqFt: double (nullable = true)
         |-- age: integer (nullable = true)
         |-- acres: double (nullable = true)
         |-- price: double (nullable = true)
         |-- features: vector (nullable = true)

        +-------+---+-----+---------+--------------------+
        |   sqFt|age|acres|    price|            features|
        +-------+---+-----+---------+--------------------+
        | 1634.0| 33| 0.04| 119900.0|  [1634.0,33.0,0.04]|
        |13837.0| 17|14.46|3500000.0|[13837.0,17.0,14.46]|
        | 9040.0| 12| 0.74|2690000.0|  [9040.0,12.0,0.74]|
         */


        //TODO 3.数据集划分0.8训练集/0.2测试集
        Dataset<Row>[] arr = vectorDF.randomSplit(new double[]{0.8, 0.2}, 100);
        Dataset<Row> trainSet = arr[0];
        Dataset<Row> testSet = arr[1];

        //TODO 4.构建线性回归模型并使用训练集训练
        LinearRegressionModel model = new LinearRegression()
                .setFeaturesCol("features")//设置特征列(应该设置向量化之后的)
                .setLabelCol("price")//设置标签列(数据中已经标记好的原本的价格)
                .setPredictionCol("predict_price")//设置预测列(后续做预测时预测的价格)
                .setMaxIter(10)//最大迭代次数
                .fit(trainSet);//使用训练集进行训练

        //TODO 5.使用测试集对模型进行测试/预测
        Dataset<Row> testResult = model.transform(testSet);
        testResult.show(false);

        //TODO 6.计算误差rmse均方误差
        double rmse = new RegressionEvaluator()//创建误差评估器
                .setMetricName("rmse") //设置要计算的误差名称,均方根误差 (sum((y-y')^2)/n)^0.5
                .setLabelCol("price")//设置真实值是哪一列
                .setPredictionCol("predict_price")//设置预测值是哪一列
                .evaluate(testResult);//对数据中的真实值和预测值进行误差计算
        System.out.println("rmse为:" + rmse);

        //TODO 7.模型保存(save)方便后续使用(load)
        //model.save("path");
        //LinearRegressionModel lmodel = LinearRegressionModel.load("path");

        //TODO 8.关闭资源
        spark.stop();
    }
}
